// pc_udp_receiver.cpp — UDP 重组 -> Draco 解码 -> PointCloud2 发布（修复悬空引用、改用“最近活跃”超时、健壮性增强）
//
// 关键变更：
// 1) 绝不在可能 erase() 之后继续使用对 frames_[seq] 的引用/迭代器；GC 放到当前帧处理完之后，并跳过当前 seq。
// 2) 超时依据 t_last（最近一次收到该帧任一分片的时间），而非首分片时间，避免长帧被误判超时。
// 3) 严格边界与一致性检查（total_chunks 上限与一致性、idx 越界防护、迟到/异常分片丢弃）。
// 4) 颜色按你“会卡死那一版”的写法：对 "rgb" 字段使用 PointCloud2Iterator<uint8_t> 逐字节写入（不再打包 float）。
// 5) 为 recvfrom 设置 SO_RCVTIMEO（200ms），析构时不再阻塞。

#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/point_cloud2_iterator.h>

#include <arpa/inet.h>
#include <sys/socket.h>
#include <unistd.h>
#include <endian.h>
#include <netinet/in.h>

#include <unordered_map>
#include <vector>
#include <string>
#include <cstring>
#include <algorithm>
#include <chrono>
#include <thread>
#include <array>
#include <atomic>
#include <limits>

// Draco
#include <draco/core/decoder_buffer.h>
#include <draco/compression/decode.h>
#include <draco/point_cloud/point_cloud.h>
#include <draco/attributes/geometry_attribute.h>

#pragma pack(push, 1)
struct UdpHeader {
  uint32_t magic;         // 'PCUD' (0x50435544)
  uint32_t seq;
  uint16_t total_chunks;
  uint16_t idx;
  uint32_t payload_size;  // size of THIS chunk's payload
  uint64_t stamp_ns;
  uint8_t  codec;         // 0=raw,1=zstd,2=draco
  uint8_t  reserved[3];
};
#pragma pack(pop)

struct FrameBuf {
  uint16_t total{0};
  std::vector<bool> got;
  std::vector<std::vector<uint8_t>> chunks;
  ros::Time stamp;
  std::chrono::steady_clock::time_point t_last;  // 最近一次收到该帧任何分片的时间
  size_t bytes_acc{0};
};

class PcReceiverDraco {
public:
  explicit PcReceiverDraco(ros::NodeHandle& nh_priv) : nh_priv_(nh_priv) {
    nh_priv_.param<std::string>("out_cloud_topic", out_topic_, std::string("/rx_cloud_xyzrgb"));
    nh_priv_.param<std::string>("bind_ip", bind_ip_, std::string("0.0.0.0"));
    nh_priv_.param<int>("bind_port", bind_port_, 50000);
    nh_priv_.param<int>("rcvbuf_mb", rcvbuf_mb_, 128);
    nh_priv_.param<double>("reassembly_timeout_s", reasm_timeout_s_, 1.2);
    nh_priv_.param<std::string>("frame_id", frame_id_, std::string("map"));

    pub_ = nh_priv_.advertise<sensor_msgs::PointCloud2>(out_topic_, 1);

    sock_ = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock_ < 0) { perror("socket"); throw std::runtime_error("socket() failed"); }

    // 允许快速复用地址
    int yes = 1;
    setsockopt(sock_, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));

    // 接收缓冲
    int rcvbuf = rcvbuf_mb_ * 1024 * 1024;
    setsockopt(sock_, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof(rcvbuf));

    // 设置接收超时，避免析构时阻塞
    timeval tv{};
    tv.tv_sec = 0;
    tv.tv_usec = 200 * 1000; // 200ms
    setsockopt(sock_, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

    std::memset(&addr_, 0, sizeof(addr_));
    addr_.sin_family = AF_INET;
    addr_.sin_port   = htons(bind_port_);
    addr_.sin_addr.s_addr = (bind_ip_ == "0.0.0.0") ? INADDR_ANY : inet_addr(bind_ip_.c_str());

    if (bind(sock_, reinterpret_cast<sockaddr*>(&addr_), sizeof(addr_)) < 0) {
      perror("bind"); throw std::runtime_error("bind() failed");
    }

    ROS_INFO_STREAM("pc_udp_receiver_draco bind " << bind_ip_ << ":" << bind_port_
                    << " out=" << out_topic_ << " frame_id=" << frame_id_
                    << " timeout_s=" << reasm_timeout_s_
                    << " rcvbuf_mb=" << rcvbuf_mb_);

    recv_thread_ = std::thread([this](){ this->recvLoop(); });
  }

  ~PcReceiverDraco(){
    running_ = false;
    if (recv_thread_.joinable()) recv_thread_.join();
    if (sock_ >= 0) close(sock_);
  }

private:
  void recvLoop() {
    std::vector<uint8_t> buf(65536);
    const uint32_t MAGIC = 0x50435544; // 'PCUD'
    const uint16_t MAX_CHUNKS = 2048;  // 安全上限，足以覆盖你日志中的 400~600 片/帧

    double last_rx_log = 0.0;

    while (running_ && ros::ok()) {
      sockaddr_in src{}; socklen_t slen = sizeof(src);
      ssize_t n = recvfrom(sock_, buf.data(), buf.size(), 0,
                           reinterpret_cast<sockaddr*>(&src), &slen);
      if (n < 0) {
        // 超时/中断，检查退出条件
        if (!running_ || !ros::ok()) break;
        continue;
      }
      if ((size_t)n < sizeof(UdpHeader)) continue;

      UdpHeader h{};
      std::memcpy(&h, buf.data(), sizeof(UdpHeader));
      if (ntohl(h.magic) != MAGIC) continue;

      const uint32_t seq   = ntohl(h.seq);
      const uint16_t total = ntohs(h.total_chunks);
      const uint16_t idx   = ntohs(h.idx);
      const uint32_t sz    = ntohl(h.payload_size);
      const uint8_t  codec = h.codec;

      // 周期性轻量日志
      const double tnow = ros::Time::now().toSec();
      if (tnow - last_rx_log > 1.0) {
        char srcip[64]; inet_ntop(AF_INET, &src.sin_addr, srcip, sizeof(srcip));
        ROS_INFO("[RX-UDP] from=%s len=%zd seq=%u idx=%u/%u payload=%u codec=%u",
                 srcip, n, seq, idx, total, sz, codec);
        last_rx_log = tnow;
      }

      // 只接受 Draco
      if (codec != 2) continue;

      // 长度一致性（单分片）
      if (sizeof(UdpHeader) + sz != static_cast<size_t>(n)) {
        ROS_WARN("[RX-UDP] bad pkt size for seq=%u idx=%u: header+payload=%zu, recv=%zd",
                 seq, idx, sizeof(UdpHeader)+static_cast<size_t>(sz), n);
        continue;
      }

      // 基本边界与上限
      if (total == 0 || total > MAX_CHUNKS) {
        ROS_WARN("[RX-UDP] bad total_chunks=%u for seq=%u, drop packet", total, seq);
        continue;
      }
      if (idx >= total) {
        ROS_WARN("[RX-UDP] bad idx=%u/%u for seq=%u, drop packet", idx, total, seq);
        continue;
      }

      // —— 获取/初始化帧缓存（此处不做 GC，避免把当前帧删掉）——
      auto it = frames_.find(seq);
      if (it == frames_.end()) {
        FrameBuf fb;
        fb.total = total;
        fb.got.assign(total, false);
        fb.chunks.resize(total);
        fb.stamp.fromNSec(be64toh(h.stamp_ns));
        fb.t_last = std::chrono::steady_clock::now();
        fb.bytes_acc = 0;
        it = frames_.emplace(seq, std::move(fb)).first;
      } else {
        // 头里的 total 如与已知不一致，直接丢弃整个旧帧，防混帧
        if (it->second.total != total) {
          ROS_WARN("[RX-UDP] total mismatch for seq=%u (%u vs %u), drop frame",
                   seq, it->second.total, total);
          frames_.erase(it);
          // 不重建同 seq 的新帧，直接跳过这个分片
          continue;
        }
      }

      FrameBuf &fb = it->second;

      // 记录分片
      if (!fb.got[idx]) {
        fb.chunks[idx].assign(buf.data()+sizeof(UdpHeader),
                              buf.data()+sizeof(UdpHeader)+sz);
        fb.got[idx] = true;
        fb.bytes_acc += sz;
        fb.t_last = std::chrono::steady_clock::now();
      }

      // 完整就发布并清理
      bool complete = std::all_of(fb.got.begin(), fb.got.end(), [](bool v){return v;});
      if (complete) {
        std::vector<uint8_t> payload;
        payload.reserve(fb.bytes_acc);
        for (uint16_t i=0;i<fb.total;++i)
          payload.insert(payload.end(), fb.chunks[i].begin(), fb.chunks[i].end());
        const ros::Time stamp = fb.stamp;
        frames_.erase(it);  // 用迭代器版本的 erase
        ROS_INFO("[RX-UDP] seq=%u complete, payload bytes=%zu", seq, payload.size());
        publishDecoded(payload, stamp);
      }

      // —— 现在再做 GC，并且显式跳过当前 seq —— 
      reassemblyGC(/*skip_seq=*/seq);
    }
  }

  void reassemblyGC(uint32_t skip_seq = std::numeric_limits<uint32_t>::max()) {
    const auto now = std::chrono::steady_clock::now();
    for (auto it = frames_.begin(); it != frames_.end(); ) {
      if (it->first == skip_seq) { ++it; continue; }
      const double idle = std::chrono::duration<double>(now - it->second.t_last).count();
      if (idle > reasm_timeout_s_) {
        const size_t got_cnt = std::count(it->second.got.begin(), it->second.got.end(), true);
        ROS_WARN("[RX-UDP] reassembly timeout for seq=%u; got=%zu/%u bytes=%zu",
                 it->first, got_cnt, it->second.total, it->second.bytes_acc);
        it = frames_.erase(it);
      } else {
        ++it;
      }
    }

    // 限制在途帧数量，避免内存爆涨
    const size_t MAX_INFLIGHT = 128;
    while (frames_.size() > MAX_INFLIGHT) {
      auto victim = std::min_element(
        frames_.begin(), frames_.end(),
        [](const auto &a, const auto &b){ return a.second.t_last < b.second.t_last; });
      ROS_WARN("[RX-UDP] too many inflight frames (%zu), drop oldest seq=%u",
               frames_.size(), victim->first);
      frames_.erase(victim);
    }
  }

  void publishDecoded(const std::vector<uint8_t>& payload, const ros::Time& stamp) {
    if (payload.size() <= 32) {
      std::string hex;
      for (size_t i=0;i<payload.size();++i){
        char b[4]; snprintf(b, sizeof(b), "%02X", payload[i]);
        hex += b; if (i+1<payload.size()) hex += " ";
      }
      ROS_WARN("RX tiny draco payload=%zu bytes, head=%s", payload.size(), hex.c_str());
    }

    ROS_INFO("[RX-Draco] decoding payload bytes=%zu stamp=%.6f",
             payload.size(), stamp.toSec());

    draco::DecoderBuffer dbuf;
    dbuf.Init(reinterpret_cast<const char*>(payload.data()), payload.size());
    draco::Decoder decoder;
    auto res = decoder.DecodePointCloudFromBuffer(&dbuf);
    if (!res.ok()) {
      ROS_WARN("Draco decode failed: %s", res.status().error_msg());
      return;
    }
    std::unique_ptr<draco::PointCloud> pc = std::move(res).value();
    const int N = pc->num_points();
    if (N <= 0) {
      ROS_WARN("Draco decoded point count <= 0");
      return;
    }

    const draco::PointAttribute* pos_att =
        pc->GetNamedAttribute(draco::GeometryAttribute::POSITION);
    const draco::PointAttribute* col_att =
        pc->GetNamedAttribute(draco::GeometryAttribute::COLOR);

    sensor_msgs::PointCloud2 msg;
    msg.header.stamp = stamp;
    msg.header.frame_id = frame_id_;
    msg.height = 1;
    msg.width  = N;

    sensor_msgs::PointCloud2Modifier mod(msg);
    if (col_att) mod.setPointCloud2FieldsByString(2, "xyz", "rgb");
    else         mod.setPointCloud2FieldsByString(1, "xyz");
    mod.resize(N);

    sensor_msgs::PointCloud2Iterator<float> it_x(msg, "x");
    sensor_msgs::PointCloud2Iterator<float> it_y(msg, "y");
    sensor_msgs::PointCloud2Iterator<float> it_z(msg, "z");
    // ===== 按旧版方式：rgb 使用 uint8_t 逐字节写入 =====
    std::unique_ptr<sensor_msgs::PointCloud2Iterator<uint8_t>> it_rgb;
    if (col_att) it_rgb.reset(new sensor_msgs::PointCloud2Iterator<uint8_t>(msg, "rgb"));

    draco::PointIndex pi(0);
    for (int i=0; i<N; ++i, ++it_x, ++it_y, ++it_z, pi = draco::PointIndex(i+1)) {
      if (pos_att) {
        const auto avi = pos_att->mapped_index(pi);
        std::array<float, 3> xyz{};
        pos_att->GetValue(avi, xyz.data());
        *it_x = xyz[0]; *it_y = xyz[1]; *it_z = xyz[2];
      } else {
        *it_x = *it_y = *it_z = 0.f;
      }

      if (col_att && it_rgb) {
        const auto avi = col_att->mapped_index(pi);
        const int comps = col_att->num_components();
        if (comps >= 4) {
          std::array<uint8_t, 4> c4{};
          col_att->GetValue(avi, c4.data());
          (*it_rgb)[0] = c4[0]; (*it_rgb)[1] = c4[1]; (*it_rgb)[2] = c4[2];
        } else if (comps == 3) {
          std::array<uint8_t, 3> c3{};
          col_att->GetValue(avi, c3.data());
          (*it_rgb)[0] = c3[0]; (*it_rgb)[1] = c3[1]; (*it_rgb)[2] = c3[2];
        } else {
          (*it_rgb)[0] = (*it_rgb)[1] = (*it_rgb)[2] = 0;
        }
        ++(*it_rgb);
      }
    }

    pub_.publish(msg);
    ROS_INFO_THROTTLE(1.0, "[RX-Draco] published pts=%d (rgb=%s)", N, col_att?"yes":"no");
  }

  ros::NodeHandle nh_priv_;
  ros::Publisher pub_;
  std::string out_topic_, bind_ip_, frame_id_;
  int bind_port_{50000}, rcvbuf_mb_{128};
  double reasm_timeout_s_{1.2};

  int sock_{-1};
  sockaddr_in addr_{};
  std::thread recv_thread_;
  std::atomic<bool> running_{true};

  // 仅 recv 线程访问
  std::unordered_map<uint32_t, FrameBuf> frames_;
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "pc_udp_receiver_draco");
  ros::NodeHandle nh_priv("~");
  try {
    PcReceiverDraco r(nh_priv);
    ros::spin();
  } catch (const std::exception& e) {
    ROS_FATAL("receiver init failed: %s", e.what());
    return 1;
  }
  return 0;
}

